from dataclasses import dataclass, field
from typing import List, Optional, Dict, Union, Any

@dataclass
class ASTNode:
    pass

@dataclass
class Program(ASTNode):
    definitions: List[ASTNode] = field(default_factory=list)

# --- 定义 ---

@dataclass
class FunctionDef(ASTNode):
    name: str
    params: List['Param']
    return_type: Optional['TypeNode']
    body: 'Block'

@dataclass
class Param(ASTNode):
    name: str
    type_annotation: Optional['TypeNode']

@dataclass
class VariableDef(ASTNode):
    name: str
    value: 'Expr'
    type_annotation: Optional['TypeNode'] = None
    is_mutable: bool = False # 默认为不可变，如规范

@dataclass
class RecordDef(ASTNode):
    """记录类型定义 (积类型)"""
    name: str
    fields: List['Param']  # 字段列表，重用Param (name, type)

# --- 表达式 ---

@dataclass
class Expr(ASTNode):
    inferred_type: Optional['TypeNode'] = field(default=None, init=False)

@dataclass
class Literal(Expr):
    value: Any
    type_name: str # '数', '词', '两仪', etc.

@dataclass
class Identifier(Expr):
    name: str

@dataclass
class BinOp(Expr):
    left: Expr
    op: str
    right: Expr

@dataclass
class UnaryOp(Expr):
    op: str
    operand: Expr

@dataclass
class CallExpr(Expr):
    func: Expr
    args: List[Expr]

@dataclass
class IfExpr(Expr):
    condition: Expr
    then_branch: 'Block'
    else_branch: Optional[Union['Block', 'IfExpr']] = None

@dataclass
class Block(Expr):
    statements: List[ASTNode]
    # 在墨算中，块中的最后一个表达式是返回值
    # 我们将通过将最后一个语句视为表达式来处理

@dataclass
class ListLiteral(Expr):
    elements: List[Expr]

@dataclass
class TupleLiteral(Expr):
    elements: List[Expr]

@dataclass
class RecordLiteral(Expr):
    """记录字面量: 点 { 横 = 3, 纵 = 4 }"""
    record_type: str  # 记录类型名称
    field_values: Dict[str, Expr]  # 字段名 -> 值

@dataclass
class FieldAccess(Expr):
    """字段访问: 点.横"""
    record: Expr  # 记录表达式
    field_name: str  # 字段名
    record_type_name: Optional[str] = field(default=None, init=False)  # 记录类型名（在语义分析时设置）



# --- 类型 ---

@dataclass
class TypeNode(ASTNode):
    name: str
    generics: List['TypeNode']

@dataclass
class FunctionType(TypeNode):
    param_types: List[TypeNode]
    return_type: TypeNode

@dataclass
class RecordType(TypeNode):
    fields: Dict[str, TypeNode] = field(default_factory=dict)
